from pathlib import Path
from time import time

def timer(func):
    def func_wrapper(*args, **kwargs):
        time_start = time()
        result = func(*args, **kwargs)
        time_end = time()
        time_spend = time_end - time_start
        print(f'\n{func.__name__} cost time {time_spend:.4f} s\n')
        return result
    return func_wrapper

def get_directory_files(directory: str):
    """
    获取目录中所有文件的名称和扩展名，以字典形式返回。
    :param directory: 目录路径
    :return: 文件名与扩展名的字典
    """
    files_dict = {}
    path = Path(directory)  # 将目录路径转换为 Path 对象
    if not path.exists():
        print(f"目录未找到: {directory}")
        return files_dict
    if not path.is_dir():
        print(f"提供的路径不是一个目录: {directory}")
        return files_dict
    
    # 使用 rglob('*') 遍历目录及其子目录中的所有文件,iterdir() 仅遍历当前目录下所有文件
    for file in path.rglob('*'):
        if file.is_file():  # 只处理文件
            name, ext = file.stem, file.suffix.lower()  # 获取文件名和扩展名
            files_dict.setdefault(name, []).append(ext)

    return files_dict

@timer
def find_missing_files(files_dict):
    """
    检查字典中具有单一扩展名的文件，判断是否缺失对应的视频或字幕文件。
    :param files_dict: 由文件名和扩展名组成的字典
    """
    video_extensions = {"mp4", "avi", "mkv"}  # 使用集合提高查找效率
    subtitle_extensions = {"srt", "ass", "vtt"}
    missing_files = []  # 列表用于存储缺失的文件信息

    for name, extensions in files_dict.items():
        if len(extensions) == 1:
            ext = extensions[0].lstrip('.')  # 去掉扩展名前的点
            if ext in video_extensions:
                missing_files.append(f"缺失字幕： {name}.{ext}")
            elif ext in subtitle_extensions:
                missing_files.append(f"缺失视频： {name}.{ext}")

    # 输出结果
    if missing_files:
        for missing in missing_files:
            print(missing)
    else:
        print("所有文件均正常，没有缺失的文件。")

if __name__ == '__main__':
    directory = r"D:\Documents\videos\4K.DEMO"
    files_dict = get_directory_files(directory)
    find_missing_files(files_dict)
